import subprocess
import threading
import sys
import time

# --- 配置部分 ---
# Go 可执行文件的路径
# 请根据你的实际情况修改，如果与脚本在同一目录，可以直接使用 "./ssh_exe.exe"
EXE_PATH = "./ssh_exe.exe"

# SSH 服务器信息 (请根据你的实际情况修改)
SSH_HOST = "192.168.31.233"
SSH_PORT = 22
SSH_USER = "zyu"
SSH_PASS = "6"

# 定义一个特殊的字符串，用于标记命令输出的结束
OUTPUT_DELIMITER = "[END_OF_OUTPUT]"

def read_pipe_output(pipe, buffer):
    """
    一个线程函数，用于在后台持续读取管道输出，直到遇到分隔符。
    """
    while True:
        try:
            line = pipe.readline().decode('utf-8')
            if not line:
                break
            buffer.append(line)
            if OUTPUT_DELIMITER in line:
                break
        except Exception as e:
            print(f"读取管道时出错: {e}", file=sys.stderr)
            break

def run_ssh_test():
    """
    启动 Go EXE 并与其进行交互测试。
    """
    print("--- 1. 启动 Go 可执行文件并建立通信管道 ---")
    try:
        # 使用 subprocess 启动 exe 并捕获标准输入/输出
        # stdin=subprocess.PIPE 允许我们向进程的标准输入写入
        # stdout=subprocess.PIPE 允许我们从进程的标准输出读取
        # stderr=subprocess.STDOUT 将 stderr 合并到 stdout
        process = subprocess.Popen(
            [EXE_PATH], 
            stdin=subprocess.PIPE, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.STDOUT,
            bufsize=1
        )
        print(f"✅ 成功启动 {EXE_PATH}")
    except FileNotFoundError:
        print(f"❌ 错误: 未找到可执行文件，请确保 {EXE_PATH} 存在。")
        return
    except Exception as e:
        print(f"❌ 启动进程时出错: {e}")
        return
    
    # 从 Go 程序读取初始的“Ready”信息
    try:
        initial_output = process.stdout.readline().decode('utf-8').strip()
        print(f"Go程序初始输出: {initial_output}")
        initial_output_2 = process.stdout.readline().decode('utf-8').strip()
        print(f"Go程序初始输出: {initial_output_2}")
    except Exception as e:
        print(f"❌ 读取Go程序初始输出时出错: {e}")
        process.terminate()
        return
    
    # --- 2. 尝试连接 SSH 服务器 ---
    print("\n--- 2. 连接 SSH 服务器 ---")
    connect_command = f"CONNECT {SSH_HOST} {SSH_PORT} {SSH_USER} {SSH_PASS}\n"
    process.stdin.write(connect_command.encode('utf-8'))
    process.stdin.flush()
    
    # 读取连接命令的响应
    connect_output_buffer = []
    read_thread = threading.Thread(target=read_pipe_output, args=(process.stdout, connect_output_buffer))
    read_thread.start()
    read_thread.join(timeout=10) # 设置一个读取超时
    
    connect_output = "".join(connect_output_buffer)
    print("连接响应:")
    print(connect_output)
    
    if "SUCCESS: Connected" not in connect_output:
        print("❌ 连接失败，测试中止。")
        process.terminate()
        return

    # --- 3. 逐一发送多条命令 ---
    commands = [
        "ls -l",
        "pwd",
        "whoami",
        "uname -a"
    ]
    
    for cmd in commands:
        print(f"\n--- 发送命令: {cmd} ---")
        command_line = f"COMMAND {cmd}\n"
        process.stdin.write(command_line.encode('utf-8'))
        process.stdin.flush()
        
        # 读取命令执行结果
        command_output_buffer = []
        read_thread = threading.Thread(target=read_pipe_output, args=(process.stdout, command_output_buffer))
        read_thread.start()
        read_thread.join(timeout=10)
        
        output = "".join(command_output_buffer).replace(f"{OUTPUT_DELIMITER}\n", "").strip()
        
        # 检查是否包含错误
        if "ERROR" in output:
            print(f"❌ 命令 '{cmd}' 执行失败，输出包含 ERROR。")
            print("---------------------------------------")
            print(output)
            print("---------------------------------------")
            # 即使失败，也继续测试下一条命令
            continue
        
        print("✅ 命令执行成功。输出:")
        print("---------------------------------------")
        print(output)
        print("---------------------------------------")

    # --- 4. 断开连接并退出程序 ---
    print("\n--- 4. 断开连接并退出 ---")
    disconnect_command = "EXIT\n"
    process.stdin.write(disconnect_command.encode('utf-8'))
    process.stdin.flush()
    
    # 关闭管道以确保进程收到结束信号
    process.stdin.close()
    
    # 等待进程终止并获取退出码
    process.wait()
    print("✅ 已发送 EXIT 命令并等待 Go 程序退出。")
    print(f"Go程序退出码: {process.returncode}")

# --- 运行测试 ---
if __name__ == "__main__":
    run_ssh_test()